import os
from tkinter import NO
from gradio.layouts import row
import pandas as pd
import gradio as gr

import plotly.express as px
from plotly.graph_objs import Figure
import plotly.graph_objects as go
from utils import EpisodeItem, TransactionFlowAgentItem



def plot_waterfall(se, title):
    se_dic = {k:v for k,v in se.items() if v != 0}
    labels = [k for k,_ in se_dic.items()]
    values = [v for _,v in se_dic.items()]
    # 定义瀑布图的数据
    # values = [1000, 200, -50, 300, -150, 250]
    # labels = ['初始值', '收入1', '支出1', '收入2', '支出2', '收入3']
    # 定义每个步骤的增减情况
    measure = ['absolute'] + ['relative'] * (len(values) - 1)

    # 创建瀑布图
    fig = go.Figure(go.Waterfall(
        name="20", orientation="v",
        measure=measure,
        x=labels,
        y=values,
        connector={"line": {"color": "rgb(63, 63, 63)"}},
    ))

    # 设置图表的标题和布局
    fig.update_layout(
        title=title,
        showlegend=True
    )

    return fig


class ShowData:
    def __init__(self, firm_id):
        self.path = "../"
        # epitem 为交易细项
        self.ep_item = EpisodeItem()
        self.tf_item = TransactionFlowAgentItem() # tf item 为交易大项目（资金流量项目）
        self.firm_id = firm_id
        self.bs_table = pd.read_excel(f"../log/{self.firm_id}_资产负债表.xlsx", sheet_name=None, 
                                      index_col=0)
        self.tf_table = pd.read_excel(f"../log/{self.firm_id}_资金流量表.xlsx", sheet_name=None, 
                                      index_col=0)

        self.bs_data_ = {k: self.transfer_data(v) for k, v in self.bs_table.items()}
        self.tf_data_ = {k: self.transfer_data(v) for k, v in self.tf_table.items()}
        self.all_info = pd.read_excel(f"../log/{self.firm_id}df_info_all.xlsx", sheet_name=None,
                                      index_col=0)
        self.periods = list(self.bs_table.keys())
        self.int_periods = [int(s) for s in self.periods]
        self.select_periods = self.periods[0]
        self.bs = self.bs_table[self.select_periods]
        self.tf = self.tf_table[self.select_periods]
        self.bs_items = self.get_bs_items()
        
    
    def get_bs_items(self):
        bs_df = self.bs_table[self.select_periods]
        names = []
        for idx in bs_df.index:
            for col in bs_df.columns[1:]:
                names.append(f"{idx}_{col}")
        return names

    @staticmethod    
    def transfer_data(df:pd.DataFrame):
        """
        把数据表的行索引名称作为第一列
        """
        df.insert(0, 'Index', df.index)  # 添加索引列，将行索引名称作为第一列
        return df

    def update_bs_tf_tables(self, year_ratio):
        """
        根据所选的年份更新资产负债表和资金流量表。
        """
        selected_year = str(year_ratio)
        bs_data = self.bs_data_[selected_year].round(1)
        tf_data = self.tf_data_[selected_year].round(1)
        ca = tf_data.loc[:,'Con.CA']
        ka = tf_data.loc[:,'Con.KA']
        aa = ca + ka
        figca = plot_waterfall(ca, '经营账户资金流量分析')
        figka = plot_waterfall(ka, '资本账户资金流量分析')
        figaa = plot_waterfall(aa, '资金流入流出分析')

        fig = px.bar(bs_data.iloc[:-1,1:])


        return bs_data, tf_data, figca, figka, figaa, fig

    def plot_bs_se(self, idx_col):
        s = idx_col.split('_')
        idx = s[0]
        col = s[1]
        se = {k:v.loc[idx,col] for k,v in self.bs_table.items()}
        new_df = pd.DataFrame(pd.Series(se))
        new_df.columns = [col]
        print(new_df)
        fig = px.line(new_df, x=new_df.index, y=new_df.columns)
        return fig

    def plot_tf_item(self, col):
        print(col, 'in?', self.tf.columns)
        dic = {k:v.loc[:,col] for k,v in self.tf_table.items()}
        df = pd.DataFrame(dic)
        df = df.loc[(df!= 0).any(axis=1)]
        df = df.T
        fig = px.line(df, x=df.index, y=df.columns)
        return fig

    def plot_tf_se(self, col):
        if col in self.all_info['info'].columns:
            df = self.all_info['info'].loc[:,[col]]
        else:
            df = self.all_info['log'].loc[:, [col]]
        fig = px.line(df, x=df.index, y=df.columns)
        return fig

    def html(self):
        with gr.Blocks(css=".dataframe tr { height: 30px; }") as demo:  # 添加 CSS 样式设置行高
            gr.Markdown("# Agent Based - Stock Flow Consistent Model(查看公司log)")
            with gr.Tab("数据全览"):
                 year_ratio = gr.Slider(
                    minimum=min(self.int_periods),
                    maximum=max(self.int_periods),
                    step=1,
                    value=min(self.int_periods),
                    label="选择年份",
                    interactive=True,
                    show_label=True,
                )
                 with gr.Row():
                    with gr.Column():
                        gr.Markdown("## 资产负债数据")
                        bs_row_names = self.bs_data_[self.periods[0]].index.tolist()
                        bs_col_names = self.bs_data_[self.periods[0]].columns.tolist()
                        bs_table = gr.Dataframe(
                            value=self.bs_data_[self.periods[0]].round(1),
                            datatype="number",
                            interactive=True,
                            col_count=(len(bs_col_names), "fixed"),
                            row_count=(len(bs_row_names), "fixed"),  # 设置行数 
                        )

                    with gr.Column():
                        gr.Markdown("## 资金流量数据")
                        tf_row_names = self.tf_data_[self.periods[0]].index.tolist()
                        tf_col_names = self.tf_data_[self.periods[0]].columns.tolist()
                        tf_table = gr.Dataframe(
                            value=self.tf_data_[self.periods[0]].round(1),
                            datatype="number",
                            interactive=True,
                            col_count=(len(tf_col_names), "fixed"),
                            row_count=(len(tf_row_names), "fixed"),  # 设置行数 
                        )
                            # 年份数据改变后，自动更新资产负债表和资金流量表    

                 ca_plot = gr.Plot()
                 ka_plot = gr.Plot()
                 aa_plot = gr.Plot()
                 bs_bs_plot = gr.Plot()

            with gr.Tab("资产负债"):
                bs_items = gr.Radio(label='资产负债', choices=self.bs_items, interactive=True)
                bs_se_plot = gr.Plot()

            with gr.Tab("资金流量"):
                info_cols = self.all_info['info'].columns.to_list()[1:]
                tf_items = gr.Radio(label='资金流量', choices=info_cols)
                tf_se_plot = gr.Plot()
            
            with gr.Tab("资金流量多维"):
                info_cols2 = ['Con.CA', 'Con.KA']
                tf_items2 = gr.Radio(label='资金流量-表',choices=info_cols2)
                tf_se_plot2 = gr.Plot()

            bs_items.change(
                fn=self.plot_bs_se,
                inputs=[bs_items],
                outputs=[bs_se_plot]
            )
            tf_items.change(
                fn=self.plot_tf_se,
                inputs=[tf_items],
                outputs=[tf_se_plot]
            )
            tf_items2.change(
                fn=self.plot_tf_item,
                inputs=[tf_items2],
                outputs=[tf_se_plot2]
            )

            year_ratio.change(
                fn=self.update_bs_tf_tables,
                inputs=[year_ratio],
                outputs=[bs_table, tf_table, ca_plot, ka_plot, aa_plot, bs_bs_plot]
            )
        return demo
    
sd = ShowData('firm_c_40')
sd.html().launch()